{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e63eb031",
   "metadata": {},
   "source": [
    "# 01 Just-In-Time Compilation\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/jit-compilation.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "1dae3ad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2fb076a",
   "metadata": {},
   "source": [
    "## How JAX transformations work\n",
    "\n",
    "The JIT compiler reduces a JAX function into a sequence of primitive operations. Primitives are units of computation; most functions in `jax.lax` represent a primitive.\n",
    "\n",
    "`jax.make_jaxpr` can show the sequence of primitives in a JAX function:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "f23ae947",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[]\u001b[39m. \u001b[34;1mlet\n",
      "    \u001b[39;22mb\u001b[35m:f32[]\u001b[39m = log a\n",
      "    c\u001b[35m:f32[]\u001b[39m = log 2.0:f32[]\n",
      "    d\u001b[35m:f32[]\u001b[39m = div b c\n",
      "  \u001b[34;1min \u001b[39;22m(d,) }\n"
     ]
    }
   ],
   "source": [
    "global_list = []\n",
    "\n",
    "\n",
    "def log2(x):\n",
    "    global_list.append(x)\n",
    "    ln_x = jnp.log(x)\n",
    "    ln_2 = jnp.log(2.0)\n",
    "    return ln_x / ln_2\n",
    "\n",
    "\n",
    "print(jax.make_jaxpr(log2)(3.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe9dc6bf",
   "metadata": {},
   "source": [
    "Notice that the side-effect of appending to the `global_list` is not captured.\n",
    "\n",
    "JAX functions must be pure functions (note: `jax.experimental.io_callback()` does allow calling side-effects at the cost of performance).\n",
    "\n",
    "When tracing a function, JAX wraps each argument in a Tracer object. The Tracer objects record all JAX operations performed on them. Then, JAX uses these records to reconstruct the entire function to produce the jaxpr.\n",
    "\n",
    "For example, since `print()` is a side-effect, any calls to it will only happen at trace-time and will not appear in the jaxpr:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "7290f952",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "printed x: JitTracer<~float32[]>\n",
      "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:f32[]\u001b[39m. \u001b[34;1mlet\n",
      "    \u001b[39;22mb\u001b[35m:f32[]\u001b[39m = log a\n",
      "    c\u001b[35m:f32[]\u001b[39m = log 2.0:f32[]\n",
      "    d\u001b[35m:f32[]\u001b[39m = div b c\n",
      "  \u001b[34;1min \u001b[39;22m(d,) }\n"
     ]
    }
   ],
   "source": [
    "def log2_with_print(x):\n",
    "    print(\"printed x:\", x)\n",
    "    ln_x = jnp.log(x)\n",
    "    ln_2 = jnp.log(2.0)\n",
    "    return ln_x / ln_2\n",
    "\n",
    "\n",
    "print(jax.make_jaxpr(log2_with_print)(3.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc64757b",
   "metadata": {},
   "source": [
    "Notice that `x` is a `Traced` object. We can also use it while debugging to print out intermediate values of a computation.\n",
    "\n",
    "JAX will only capture the function based on the execution at trace-time. For example, if we have a Python conditional, the jaxpr will only know about the branch we take:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "262626ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{ \u001b[34;1mlambda \u001b[39;22m; a\u001b[35m:i32[3]\u001b[39m. \u001b[34;1mlet\u001b[39;22m  \u001b[34;1min \u001b[39;22m(a,) }\n"
     ]
    }
   ],
   "source": [
    "def log2_if_rank_2(x):\n",
    "    if x.ndim == 2:\n",
    "        ln_x = jnp.log(x)\n",
    "        ln_2 = jnp.log(2.0)\n",
    "        return ln_x / ln_2\n",
    "    else:\n",
    "        return x\n",
    "\n",
    "\n",
    "print(jax.make_jaxpr(log2_if_rank_2)(jnp.array([1, 2, 3])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df46cc68",
   "metadata": {},
   "source": [
    "Since we took the else branch at trace-time, the jaxpr only knows to immediately return the first argument.\n",
    "\n",
    "## JIT compiling a function\n",
    "\n",
    "Example of JIT compiling a `selu`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "185a0a8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0027778148651123047\n",
      "0.048252105712890625\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "def selu(x, alpha=1.67, lambda_=1.05):\n",
    "    return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)\n",
    "\n",
    "\n",
    "x = jnp.arange(1_000_000)\n",
    "start = time.time()\n",
    "selu(x).block_until_ready()\n",
    "print(time.time() - start)\n",
    "\n",
    "start = time.time()\n",
    "jax.jit(selu)(x).block_until_ready()\n",
    "print(time.time() - start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a54d6243",
   "metadata": {},
   "source": [
    "Before we JIT compile it, we are sending one operation at a time to the accelerator, which limits the XLA compiler’s ability to optimize anything.\n",
    "\n",
    "`jax.jit()` gives the XLA compiler all of the code upfront. On the first call, the XLA compiler will trace the function and emit jaxpr which is then compiled into efficient accelerator-specific instructions.\n",
    "\n",
    "Any subsequent calls will use the efficient compiled accelerator instructions directly.\n",
    "\n",
    "## Why can’t we JIT everything?\n",
    "\n",
    "Cannot JIT functions that use Python conditionals or loops.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "8742d300",
   "metadata": {},
   "outputs": [
    {
     "ename": "TracerBoolConversionError",
     "evalue": "Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function g at /var/folders/r0/zlgwr7551fg116vdpn_9s7ph0000gn/T/ipykernel_85568/222786207.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument n.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mTracerBoolConversionError\u001b[39m                 Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[37]\u001b[39m\u001b[32m, line 9\u001b[39m\n\u001b[32m      5\u001b[39m         i += \u001b[32m1\u001b[39m\n\u001b[32m      6\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m x + i\n\u001b[32m----> \u001b[39m\u001b[32m9\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[33mAnswer:\u001b[39m\u001b[33m\"\u001b[39m, \u001b[43mg\u001b[49m\u001b[43m(\u001b[49m\u001b[32;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[32;43m20\u001b[39;49m\u001b[43m)\u001b[49m)\n",
      "    \u001b[31m[... skipping hidden 13 frame]\u001b[39m\n",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[37]\u001b[39m\u001b[32m, line 4\u001b[39m, in \u001b[36mg\u001b[39m\u001b[34m(x, n)\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;129m@jax\u001b[39m.jit\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mg\u001b[39m(x, n):\n\u001b[32m      3\u001b[39m     i = \u001b[32m0\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m4\u001b[39m     \u001b[38;5;28;01mwhile\u001b[39;00m i < n:\n\u001b[32m      5\u001b[39m         i += \u001b[32m1\u001b[39m\n\u001b[32m      6\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m x + i\n",
      "    \u001b[31m[... skipping hidden 1 frame]\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/Programming/Learn JAX/.venv/lib/python3.12/site-packages/jax/_src/core.py:1721\u001b[39m, in \u001b[36mconcretization_function_error.<locals>.error\u001b[39m\u001b[34m(self, arg)\u001b[39m\n\u001b[32m   1720\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34merror\u001b[39m(\u001b[38;5;28mself\u001b[39m, arg):\n\u001b[32m-> \u001b[39m\u001b[32m1721\u001b[39m   \u001b[38;5;28;01mraise\u001b[39;00m TracerBoolConversionError(arg)\n",
      "\u001b[31mTracerBoolConversionError\u001b[39m: Attempted boolean conversion of traced array with shape bool[].\nThe error occurred while tracing the function g at /var/folders/r0/zlgwr7551fg116vdpn_9s7ph0000gn/T/ipykernel_85568/222786207.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument n.\nSee https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def g(x, n):\n",
    "    i = 0\n",
    "    while i < n:\n",
    "        i += 1\n",
    "    return x + i\n",
    "\n",
    "\n",
    "print(\"Answer:\", g(10, 20))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fca36fc",
   "metadata": {},
   "source": [
    "Since the sequence of operations depends on a runtime value not known at trace-time so it is not possible to compile.\n",
    "\n",
    "Traced values in JIT like `x` and `n` can only affect control-flow based on their static attributes like `shape` or `dtype`, not their values:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "194cc0d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 20\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def g(x, y):\n",
    "    if x.dtype == jax.dtypes.bfloat16:\n",
    "        return x\n",
    "    else:\n",
    "        return y\n",
    "\n",
    "\n",
    "print(\"Answer:\", g(10, 20))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39708c52",
   "metadata": {},
   "source": [
    "Generally, avoid conditionals on values. If conditionals are necessary, use something like `jax.lax.cond()` for conditionals on traced values.\n",
    "\n",
    "However, sometimes even that is not enough. In that case, we can JIT compile part of the function. For example, if the computationally expensive portion of the function is the loop body, we could JIT compile that:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37a6d0b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 30\n"
     ]
    }
   ],
   "source": [
    "@jax.jit\n",
    "def loop_body(prev_i):  # <PjitFunction of <function loop_body at 0x102c56340>>\n",
    "    return prev_i + 1\n",
    "\n",
    "\n",
    "def g(x, n):  # <function g at 0x10aa4c900>\n",
    "    i = 0\n",
    "    while i < n:\n",
    "        i = loop_body(i)\n",
    "    return x + i\n",
    "\n",
    "\n",
    "print(\"Answer:\", g(10, 20))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18e287a6",
   "metadata": {},
   "source": [
    "## Marking arguments as static\n",
    "\n",
    "If we really need to JIT compile a function that is dependent on a runtime value, we can mark the argument static which will wrap it in a less restrictive tracer.\n",
    "\n",
    "The downside is that the emitted jaxpr is specific to the value passed (the value is treated as a constant in the jaxpr). For every new value of the static argument, JAX will recompile the function.\n",
    "\n",
    "It is recommended only to use static arguments for inputs with a fixed set of possible values.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05127f8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Answer: 30\n"
     ]
    }
   ],
   "source": [
    "from functools import partial\n",
    "\n",
    "\n",
    "@partial(jax.jit, static_argnames=(\"n\"))\n",
    "def g(x, n):\n",
    "    i = 0\n",
    "    while i < n:\n",
    "        i += 1\n",
    "    return x + i\n",
    "\n",
    "\n",
    "print(\"Answer:\", g(10, 20))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96de339d",
   "metadata": {},
   "source": [
    "## JIT and caching\n",
    "\n",
    "When a function marked with `@jax.jit`, it will be compiled and the emitted XLA will be cached. Subsequent calls to the function will use the cached XLA, so the overhead is amortized out.\n",
    "\n",
    "If we specify `static_argnums` / `static_argnames`, only the XLA for the specific set of inputs to the static arguments will be cached. If any of the inputs to the static arguments changes, the function must be recompiled.\n",
    "\n",
    "Do not call `jax.jit()` on temporary functions defined inside loops. Since the cache relies on the hash of the function, it becomes problematic when equivalent functions are redefined. This will recompile the function on each temporary definition.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbec952b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.22260403633117676\n",
      "0.011249303817749023\n"
     ]
    }
   ],
   "source": [
    "def loop_body(i):\n",
    "    return i + 1\n",
    "\n",
    "\n",
    "def g_temporary_inner_function(x, n):\n",
    "    i = 0\n",
    "    while i < n:\n",
    "        i = jax.jit(lambda x: loop_body(x))(i)  # Bad! Recompilation on each iteration.\n",
    "    return x + i\n",
    "\n",
    "\n",
    "def g_inner_function(x, n):\n",
    "    i = 0\n",
    "    while i < n:\n",
    "        i = jax.jit(loop_body)(i)  # Good! Subsequent calls will hit XLA cache.\n",
    "    return x + i\n",
    "\n",
    "\n",
    "start = time.time()\n",
    "g_temporary_inner_function(10, 20).block_until_ready()\n",
    "print(time.time() - start)\n",
    "\n",
    "start2 = time.time()\n",
    "g_inner_function(10, 20).block_until_ready()\n",
    "print(time.time() - start2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90a2ce08",
   "metadata": {},
   "source": [
    "The performance of the first function is significantly worse because it recompiles to XLA after each loop iteration. The second function hits the XLA cache after the second loop iteration.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
